【DLR】SageMaker Neoでコンパイルしたモデルを動かすときに使えるTips(起動高速化・処理高速化)
カフェチームの山本です。
作成した学習モデルをSageMaker Neoを利用してコンパイルし、Jetson Xavier NXなどのシングルボードで動かす、といった場面があると思います。今回は、コンパイルしたモデルを動かす際に利用するDLR(Deep Learning Runtime)について、使える小ネタを紹介したいと思います。簡単なものですが、起動高速化・処理高速化など、利用するとかなり便利なのでおすすめです。
DLR(Deep Learning Runtime)
シングルボードなどで学習モデルを動かしたいとき、モデルそのままでは大きく動かない(もしくは処理が遅い)という場合があります。そうしたモデルはSageMaker Neoでコンパイルすることで、指定したデバイス用に最適化することができます。このモデルは、DLR(Deep Learning Runtime)というライブラリで動くようになっています。
今回は、このDLRについて、紹介されることが少ないものの、結構便利なTipsを紹介します。
SageMaker Neoでモデルをコンパイル(最適化)して動かす方法自体については、詳しくは同じカフェチームのSINの執筆した記事で詳しく解説されているため、そちらをご覧ください。
[Amazon SageMaker] イメージ分類のモデルをNeoで最適化して、Jetson Nano+OpenCV+Webカメラで使用してみました | Developers.IO
DLRで使える小ネタ
今回紹介するのは主に、以下のページの内容です
Additional Options for TensorRT Optimized Models - dlr documentation
モデルを動かすスクリプトとしては、以下のようなものを用意したとします。ファイル名はexample.pyとします。(このスクリプトである必要はなく、DLRでモデルを利用している一例です)
import dlr import numpy as np # Load model. # /path/to/model is a directory containing the compiled model artifacts (.so, .params, .json) model = dlr.DLRModel('/path/to/model', 'gpu', 0) # Prepare some input data. x = np.random.rand(1, 3, 224, 224) # Run inference. y = model.run(x)
通常、スクリプトを起動するコマンドは以下のようだと思います。
python3 example.py
この起動方法を、フラグ(Additional Flags)を追加してから起動することによって、動作を変えられます。フラグは複数同時に使用できます。
ビルドのキャッシュ(起動高速化)
DLRの仕様として、通常の起動方法だと、毎回モデルを読み込む度にビルドし直すようになっています。そのため、毎回起動するたびに待つ必要がありました。(開発中に何度も試行錯誤しようとすると、開発スピードが遅くなり、やる気も低下するという問題がありました。)
起動方法を以下のようにすることで、ビルドした内容をキャッシュし、保存することができます。これによって、1回目の起動時にはビルドが行われるため時間がかかりますが、2回目以降はキャッシュした内容がロードされるだけであるため、時間が大幅に短縮されます。
TVM_TENSORRT_CACHE_DIR=. python3 example.py
指定したフォルダに「9326145774222630086.meta」「9326145774222630086.plan」のようなファイルが出力されます。
具体的には、Jetson Xavier NXにおいて自分が利用してる骨格検出のモデル(SageMaker Neoでコンパイルしたもの)の場合、通常のビルドで2分ほどかかっていた起動が、キャッシュのロードのみでは6秒ほどに短縮されました。かなり開発が効率化されました。おすすめです。
モデルのビット数の変換(処理速度高速化)
モデルを浮動小数点16ビットに変換して実行します。上のページに「これによりパフォーマンスが大幅に向上しますが、モデルの精度が若干低下する可能性があります」と書かれているように、ビット数を落とすことで精度を少し下げながらも、速度を大幅に上げることができます。
一般的に、16bitにしてもそこまで精度が落ちず、むしろ高速化の度合いが大きいことから、特にエッジデバイスで良く利用される方法です。
コマンドとしては、以下のようです。
TVM_TENSORRT_USE_FP16=1 python3 example.py
具体的には、先程の自分が利用しているモデルの場合ですが、処理時間が140[ms]程度から60[ms]程度になり、2倍以上の高速化ができました(処理結果の精度の方は確認できていません)。
SageMaker Neoなどでモデルを再度作成する必要がなく、コマンド1つで処理を高速化できるためかなり便利です。1度は16bitで試しに実行し精度を確認してみて、問題なさそうなら採用する、という流れが良いと思います。
ワークスペースサイズの変更(スループット増加)
モデルを動かす際に利用するGPUのサイズを変更できます。層の深いモデルや、入力サイズが大きいモデルを利用する場合や、データをまとめて入力したいときに利用できるようです。
上のリンクのページに「Jetson NanoおよびJetson TX1ターゲットの場合は最大ワークスペースサイズを256メガバイトに、他のすべてのNVIDIA GPUターゲットの場合は1ギガバイトに自動的に設定します」と書かれているように、SageMaker Neoでコンパイルしたモデルは、動くボードによって自動でこの値を、もともと設定しているようです。この値を上書きして使用します。
コマンドとしては、以下のようです(2GBの場合)。
TVM_TENSORRT_MAX_WORKSPACE_SIZE=2147483647 python3 run.py
自分は今の所使う機会がなく、具体的な効果がわかっていません。
まとめ
今回は、DLRを利用する際に、コマンド1つで起動高速化・処理速度高速化・スループット増加できる方法を紹介しました。
3つのうち最初の2つは、実際に使ってみたところ大きな効果があったのでオススメです。(最後の1つはまだ効果がわかっていません。)
参考にさせていただいたページ
Additional Options for TensorRT Optimized Models - dlr documentation
[Amazon SageMaker] イメージ分類のモデルをNeoで最適化して、Jetson Nano+OpenCV+Webカメラで使用してみました | Developers.IO